#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
逆运动学重放分析器
用于重放和分析逆运动学计算的耗时
"""

import os
import sys
import time
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
import argparse
import pickle
import pandas as pd
import cProfile
import pstats
from io import StringIO
import time

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.append(project_root)
sys.path.append(current_dir)

# 导入自定义模块
from utils.ik_data_recorder import IKDataRecorder
from robot_control.fangnuo_car_ik_new import FangNuoCarIK
from robot_control.fangnuo_car import FangnuoCar
from utils.geo_math import print_pose_info

class IKReplayAnalyzer:
    def __init__(self, data_file=None, visualization=True, use_real_arms=False):
        """
        初始化重放分析器
        
        参数:
            data_file: 要分析的数据文件路径，如果为None则使用最新的
            visualization: 是否启用可视化
        """
        # 初始化数据记录器
        self.data_recorder = IKDataRecorder()
        
        # 如果未指定数据文件，使用最新的
        if data_file is None:
            data_file = self.data_recorder.get_latest_session_file()
            if data_file is None:
                print("错误：未找到数据文件")
                sys.exit(1)
        
        # 加载数据
        self.data_list = self.data_recorder.load_session(data_file)
        if not self.data_list:
            print(f"错误：无法从 {data_file} 加载数据")
            sys.exit(1)
        
        # 初始化IK求解器
        self.ik_solver = FangNuoCarIK(Unit_Test=False, Visualization=visualization)
        
        # 初始化机械臂控制器（如果启用真机控制）
        self.use_real_arms = use_real_arms
        self.arm_ctrl = None
        if self.use_real_arms:
            print("初始化真实机械臂控制器...")
            self.arm_ctrl = FangnuoCar(use_real_arms=True)
            print("真实机械臂控制器初始化完成")
        
        # 性能分析结果
        self.performance_results = []
        time.sleep(1)
        
        print(f"IK重放分析器初始化完成，加载了 {len(self.data_list)} 条数据记录，{'启用' if self.use_real_arms else '未启用'}真机控制")
    
    def run_single_replay(self, data_record, profile=False):
        """
        重放单条数据记录
        
        参数:
            data_record: 数据记录
            profile: 是否进行性能分析
            
        返回:
            耗时和结果
        """
        left_wrist = data_record['left_wrist'].copy()
        right_wrist = data_record['right_wrist'].copy()
        current_lr_arm_q = data_record['current_lr_arm_q']
        current_lr_arm_dq = data_record['current_lr_arm_dq']
        
        # 绕Y轴旋转90度的旋转矩阵
        # [ cos(90°)  0  sin(90°) ]
        # [    0      1     0     ]
        # [-sin(90°)  0  cos(90°) ]
        # 即:
        # [ 0  0  1 ]
        # [ 0  1  0 ]
        # [-1  0  0 ]
        rot_y_90 = np.array([
            [0, 0, 1],
            [0, 1, 0],
            [-1, 0, 0]
        ])
        
        # 先平移再旋转
        # 左右手臂沿着Z轴向上平移10厘米
        # left_wrist[0, 3] += 0.1  # Z轴平移0.1米
        # right_wrist[0, 3] += 0.1  # Z轴平移0.1米
        # left_wrist[2, 3] += 0.1  # Z轴平移0.1米
        # right_wrist[2, 3] += 0.1  # Z轴平移0.1米
        
        # 应用旋转到左右手臂位姿
        # left_wrist[:3, :3] = np.matmul(rot_y_90, left_wrist[:3, :3])
        # right_wrist[:3, :3] = np.matmul(rot_y_90, right_wrist[:3, :3])
        
        # 记录开始时间
        start_time = time.time()
        
        if profile:
            # 使用cProfile进行性能分析
            pr = cProfile.Profile()
            pr.enable()
            
            # 执行IK计算
            sol_q, sol_tauff = self.ik_solver.solve_ik(left_wrist, right_wrist, current_lr_arm_q, current_lr_arm_dq)
            time.sleep(0.5)
            # 停止性能分析
            pr.disable()
            
            # 获取性能分析结果
            s = StringIO()
            ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
            ps.print_stats(20)  # 打印前20个耗时最多的函数
            profile_result = s.getvalue()
        else:
            # 直接执行IK计算
            sol_q, sol_tauff = self.ik_solver.solve_ik(left_wrist, right_wrist, current_lr_arm_q, current_lr_arm_dq)
            profile_result = None
        
        # 计算耗时
        elapsed_time = time.time() - start_time
        print("-----left_wrist/n")
        print_pose_info(left_wrist)
        print("-----right_wrist/n")
        print_pose_info(right_wrist)
        print(f"sol_q left: {sol_q[:6]}")
        print(f"sol_q right: {sol_q[6:12]}")  
        print("*"*50)  
        
        # 如果启用真机控制，控制机械臂
        if self.use_real_arms and self.arm_ctrl is not None:
            # 打印真机左臂和右臂的关节位置
            robot_status = self.arm_ctrl.status()
            
            print(f"左臂关节位置: {robot_status['left_arm']['position']}")
            print(f"右臂关节位置: {robot_status['right_arm']['position']}")
           
            self.arm_ctrl.ctrl_dual_arm(sol_q, sol_tauff)
            print("机械臂控制命令已发送")
        
        return {
            'elapsed_time': elapsed_time,
            'sol_q': sol_q,
            'sol_tauff': sol_tauff,
            'profile_result': profile_result
        }
    
    def run_batch_replay(self, num_samples=None, profile_first=True):
        """
        批量重放数据
        
        参数:
            num_samples: 要重放的样本数，如果为None则使用全部
            profile_first: 是否对第一个样本进行性能分析
            
        返回:
            性能结果列表
        """
        # 确定要处理的样本数
        if num_samples is None or num_samples > len(self.data_list):
            num_samples = len(self.data_list)
        
        print(f"开始批量重放 {num_samples} 条数据记录...")
        
        # 清空性能结果
        self.performance_results = []
        
        # 处理每条数据记录
        for i, data_record in enumerate(self.data_list[:num_samples]):
            # 对第一个样本进行性能分析
            profile = (i == 0 and profile_first)
            
            if i%10 == 0:
                print(f"-------------------------- {i}")
            # 重放数据
            result = self.run_single_replay(data_record, profile)
            
            # 记录性能结果
            self.performance_results.append({
                'index': i,
                'timestamp': data_record['timestamp'],
                'elapsed_time': result['elapsed_time'],
                'profile_result': result['profile_result'] if profile else None
            })
            
            time.sleep(0.2)
            # 打印进度
            if (i + 1) % 10 == 0 or i == num_samples - 1:
                print(f"已处理 {i + 1}/{num_samples} 条记录")
        
        return self.performance_results
    
    def analyze_performance(self):
        """
        分析性能结果
        
        返回:
            性能统计信息
        """
        if not self.performance_results:
            print("错误：没有性能结果可分析")
            return None
        
        # 提取耗时数据
        elapsed_times = [result['elapsed_time'] for result in self.performance_results]
        
        # 计算统计信息
        stats = {
            'count': len(elapsed_times),
            'min': min(elapsed_times),
            'max': max(elapsed_times),
            'mean': np.mean(elapsed_times),
            'median': np.median(elapsed_times),
            'std': np.std(elapsed_times),
            'p95': np.percentile(elapsed_times, 95),
            'p99': np.percentile(elapsed_times, 99)
        }
        
        # 打印统计信息
        print("\n性能统计信息:")
        print(f"样本数: {stats['count']}")
        print(f"最小耗时: {stats['min']:.6f} 秒")
        print(f"最大耗时: {stats['max']:.6f} 秒")
        print(f"平均耗时: {stats['mean']:.6f} 秒")
        print(f"中位数耗时: {stats['median']:.6f} 秒")
        print(f"标准差: {stats['std']:.6f} 秒")
        print(f"95百分位耗时: {stats['p95']:.6f} 秒")
        print(f"99百分位耗时: {stats['p99']:.6f} 秒")
        
        # 如果有性能分析结果，打印第一个
        if self.performance_results[0]['profile_result']:
            print("\n详细性能分析 (第一个样本):")
            print(self.performance_results[0]['profile_result'])
        
        return stats
    
    def plot_performance(self, output_file=None):
        """
        绘制性能图表
        
        参数:
            output_file: 输出文件路径，如果为None则显示图表
        """
        if not self.performance_results:
            print("错误：没有性能结果可绘图")
            return
        
        # 创建DataFrame
        df = pd.DataFrame(self.performance_results)
        
        # 设置图表样式
        plt.style.use('ggplot')
        plt.figure(figsize=(12, 8))
        
        # 绘制耗时曲线
        plt.subplot(2, 1, 1)
        plt.plot(df.index, df['elapsed_time'], 'b-', linewidth=1)
        plt.title('ik replay analyzer', fontsize=16)
        plt.ylabel('elapsed time (sec)', fontsize=12)
        plt.grid(True)
        
        # 绘制耗时分布直方图
        plt.subplot(2, 1, 2)
        plt.hist(df['elapsed_time'], bins=30, alpha=0.7, color='blue')
        plt.xlabel('elapsed time (sec)', fontsize=12)
        plt.ylabel('count', fontsize=12)
        plt.title('elapsed time distribution', fontsize=14)
        plt.grid(True)
        
        plt.tight_layout()
        
        # 保存或显示图表
        if output_file:
            plt.savefig(output_file)
            print(f"性能图表已保存到: {output_file}")
        else:
            plt.show()

def main():
    """
    主函数
    """
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='逆运动学重放分析器')
    parser.add_argument('--data-file', type=str, default="./data/ik_data/ik_data_20250606_172707.pkl", help='要分析的数据文件路径')
    parser.add_argument('--num-samples', type=int, help='要分析的样本数')
    parser.add_argument('--output', type=str, help='输出图表文件路径')
    parser.add_argument('--real-arm', action='store_true', default=False, help='是否使用真实机械臂控制')
    args = parser.parse_args()
    
    # 创建分析器
    analyzer = IKReplayAnalyzer(args.data_file, use_real_arms=args.real_arm)
    
    # 运行批量重放
    analyzer.run_batch_replay(args.num_samples)
    
    # 如果使用真实机械臂，在结束时回到初始位置
    if args.real_arm and analyzer.arm_ctrl is not None:
        print("重放完成，机械臂回到初始位置...")
        analyzer.arm_ctrl.go_home()
        time.sleep(5.0)  # 等待机械臂回到初始位置
        analyzer.arm_ctrl.stop()
    
    # 分析性能
    # analyzer.analyze_performance()
    
    # 绘制性能图表
    # analyzer.plot_performance(args.output)

if __name__ == "__main__":
    main()
